package com.lzh.crypt.rsa;

import com.lzh.crypt.enums.RsaAlgorithmEnum;
import com.lzh.crypt.enums.SignAlgorithmEnum;
import com.lzh.crypt.exception.ErrorCode;
import com.lzh.crypt.exception.RSAException;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.Signature;
import java.security.interfaces.RSAPrivateKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;

/**
 * @author zhehen.lu
 * @date 2025/8/10 19:53
 */
public class RSAPrivateKeyPair {

    private RSAPrivateKey rsaPrivateKey;

    private String algorithm = "RSA";

    private String signAlgorithm = "MD5withRSA";

    public RSAPrivateKeyPair(RSAPrivateKey rsaPrivateKey) {
        this.rsaPrivateKey = rsaPrivateKey;
    }

    public RSAPrivateKeyPair(RSAPrivateKey rsaPrivateKey, SignAlgorithmEnum signAlgorithmEnum, RsaAlgorithmEnum rsaAlgorithmEnum) {
        this.rsaPrivateKey = rsaPrivateKey;
        this.signAlgorithm = signAlgorithmEnum.getSignType();
        this.algorithm = rsaAlgorithmEnum.getValue();
    }

    public RSAPrivateKeyPair(String privateKeyBase64) {
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKeyBase64));
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            this.rsaPrivateKey = (RSAPrivateKey) keyFactory.generatePrivate(pkcs8EncodedKeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RSAException(ErrorCode.RSA_PRIVATE_ERROR.getErrorMsg(),e);
        }

    }

    public RSAPrivateKeyPair(String privateKeyBase64, String algorithm, String signAlgorithm) {
        this.algorithm = algorithm;
        this.signAlgorithm = signAlgorithm;
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(privateKeyBase64));
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            this.rsaPrivateKey = (RSAPrivateKey) keyFactory.generatePrivate(pkcs8EncodedKeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new RSAException(ErrorCode.RSA_PRIVATE_ERROR.getErrorMsg(),e);
        }
    }

    /**
     * @return the algorithm
     */
    public String getAlgorithm() {
        return algorithm;
    }

    /**
     * @param algorithm the algorithm to set
     */
    public void setAlgorithm(String algorithm) {
        this.algorithm = algorithm;
    }

    /**
     * @return the signAlgorithm
     */
    public String getSignAlgorithm() {
        return signAlgorithm;
    }

    /**
     * @param signAlgorithm the signAlgorithm to set
     */
    public void setSignAlgorithm(String signAlgorithm) {
        this.signAlgorithm = signAlgorithm;
    }

    /**
     * @param rsaPrivateKey the rsaPrivateKey to set
     */
    public void setRsaPrivateKey(RSAPrivateKey rsaPrivateKey) {
        this.rsaPrivateKey = rsaPrivateKey;
    }

    public RSAPrivateKey getRsaPrivateKey() {
        return rsaPrivateKey;
    }

    public String getPrivatekeyBase64() {
        return Base64.encodeBase64String(rsaPrivateKey.getEncoded());
    }

    public byte[] sign(byte[] content) {
        try {
            Signature signature = Signature.getInstance(signAlgorithm);
            signature.initSign(rsaPrivateKey);
            signature.update(content);
            return signature.sign();
        } catch (Exception e) {
            throw new RSAException(ErrorCode.RSA_SIGNATURE_ERROR.getErrorMsg(),e);
        }
    }

    public byte[] encrypt(byte[] bs) {
        try {
            Cipher cipher = Cipher.getInstance(algorithm);
            cipher.init(Cipher.ENCRYPT_MODE, rsaPrivateKey);
            return cipher.doFinal(bs);
        } catch (Exception e) {
            throw new RSAException(ErrorCode.RSA_ENCRYPTION_EXCEPTION.getErrorMsg(),e);
        }
    }

    public byte[] decypt(byte[] bs) {
        try {
            Cipher cipher = Cipher.getInstance(algorithm);
            cipher.init(Cipher.DECRYPT_MODE, rsaPrivateKey);
            return cipher.doFinal(bs);
        } catch (Exception e) {
            throw new RSAException(ErrorCode.RSA_DECRYPTION_EXCEPTION.getErrorMsg(),e);
        }
    }

}
